import nvtx

class NVTXDataLoaderDecorator():
    def __init__(self, color='green', desc="DataLoader"):
        self.desc = desc
        self.color = color

    def __call__(self, cls):
        # 接收被装饰的类
        class Wrapper(cls):
            desc = self.desc
            color = self.color

            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.desc = self.desc or cls.__name__
            
            def __getitem__(self, idx):
                with nvtx.annotate(message=self.desc, color=self.color):
                    return super().__getitem__(idx)

        return Wrapper